from fastapi import FastAPI, File, Form, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
import os
import shutil
from typing import List
import json

from document_processor import DocumentProcessor
from vector_store import VectorStore
from qa_engine import QAEngine

app = FastAPI(title="AI文件问答助手", version="1.0.0")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化组件
doc_processor = DocumentProcessor()
vector_store = VectorStore()
qa_engine = QAEngine()

# 创建上传目录
UPLOAD_DIR = "uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)

@app.post("/api/upload")
async def upload_files(files: List[UploadFile] = File(...)):
    """上传文件并处理"""
    try:
        uploaded_files = []
        
        for file in files:
            # 保存文件
            file_path = os.path.join(UPLOAD_DIR, file.filename)
            with open(file_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)
            
            # 处理文档
            content = doc_processor.process_file(file_path)
            
            if content:
                # 添加到向量存储
                vector_store.add_documents(content, file.filename)
                uploaded_files.append(file.filename)
        
        return {
            "status": "success",
            "message": f"成功上传并处理 {len(uploaded_files)} 个文件",
            "files": uploaded_files
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/ask")
async def ask_question(question: str = Form(...)):
    """提问并获取答案"""
    try:
        # 检索相关文档
        relevant_docs = vector_store.search(question, k=3)
        
        if not relevant_docs:
            return {
                "answer": "抱歉，没有找到相关的文档内容。",
                "sources": []
            }
        
        # 生成答案
        answer = qa_engine.generate_answer(question, relevant_docs)
        
        # 格式化来源信息
        sources = []
        for doc in relevant_docs:
            sources.append({
                "file": doc["metadata"]["filename"],
                "content": doc["content"][:200] + "..." if len(doc["content"]) > 200 else doc["content"],
                "score": doc["score"]
            })
        
        return {
            "answer": answer,
            "sources": sources
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/status")
async def get_status():
    """获取系统状态"""
    return {
        "documents_count": vector_store.get_document_count(),
        "status": "running"
    }

@app.delete("/api/clear")
async def clear_all():
    """清除所有文档"""
    try:
        vector_store.clear()
        # 清除上传的文件
        for filename in os.listdir(UPLOAD_DIR):
            file_path = os.path.join(UPLOAD_DIR, filename)
            if os.path.isfile(file_path):
                os.remove(file_path)
        
        return {"status": "success", "message": "已清除所有文档"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
